import socket
import threading
import time
import sys
import base64
import face_recognition
import re
import os
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad

def PicEncrypt(PicPath):
    password = b'ctseu11451419198' 
    
    iv = b'qwqqvq1231321321'
    print(len(iv))
    with open(PicPath,'rb') as f:
        pic = f.read()
    PadPic = pad(pic, AES.block_size)
    aes = AES.new(password,AES.MODE_CBC,iv)
    en_text = aes.encrypt(PadPic) 
    with open(PicPath + '.aes','wb+') as f:
        f.write(en_text)
    

def PicDecrypt(PicPath):
    password = b'ctseu11451419198' 
    
    iv = b'qwqqvq1231321321'
    print(len(iv))
    with open(PicPath + '.aes','rb') as f:
        pic = f.read()
    aes = AES.new(password,AES.MODE_CBC,iv) 
    DePic = aes.decrypt(pic)

    OriPic = unpad(DePic, AES.block_size)
    with open(PicPath,'wb+') as f:
        f.write(OriPic)

def myrecongnition():
    PicDecrypt('know/lx1.jpg')
    PicDecrypt('know/lbw1.jpg')
    lbw_image = face_recognition.load_image_file("./know/lbw1.jpg")
    lx_image = face_recognition.load_image_file("./know/lx1.jpg")
    unknown_image = face_recognition.load_image_file("./data/test.jpg")

    try:
        lbw_face_encoding = face_recognition.face_encodings(lbw_image)[0]
        lx_face_encoding = face_recognition.face_encodings(lx_image)[0]
        unknown_face_encoding = face_recognition.face_encodings(unknown_image)[0]
    except IndexError:
        print("I wasn't able to locate any faces in at least one of the images. Check the image files. Aborting...")
        os.system('rm ./know/lbw1.jpg')
        os.system('rm ./know/lx1.jpg')
        return ['Unable to check the face']

    os.system('rm ./know/lbw1.jpg')
    os.system('rm ./know/lx1.jpg')

    name = ['Lu Benwei', 'Luo Xiang']

    known_faces = [
        lbw_face_encoding,
        lx_face_encoding
    ]

    results = face_recognition.compare_faces(known_faces, unknown_face_encoding)

    my_result = []
    for i in range(len(results)):
        if results[i]:
            my_result.append(name[i])
        else:
            my_result.append('Unknow')
    
    return my_result

def socket_service():
    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        s.bind(('172.17.53.112', 5555))
        s.listen(10)
    except socket.error as msg:
        print(msg)
        sys.exit(1)
    print('Waiting connection...')
 
    while 1:
        conn, addr = s.accept()
        t = threading.Thread(target=deal_data, args=(conn, addr))
        t.start()
 
def deal_data(conn, addr):
    print('Accept new connection from {0}'.format(addr))
    conn.send(('Hi, Welcome to the server!').encode())
    while 1:
        data = conn.recv(1024)
        print(data)
        print('{0} client send data'.format(addr)) 
        time.sleep(1)
        data_decode = data.decode('ascii')
        if data_decode[:4] == 'imag':
            print('ok')
            print(data_decode)
            img_size = int(re.search(r'imag, size:(\d+)', data_decode).group(1))
            print('size: ' + str(img_size))
            img_buff = b''
            recvd_size = 0
            conn.send(bytes('OK',"ascii"))

            while not recvd_size == img_size:
                if img_size - recvd_size > 1024:
                    data = conn.recv(1024)
                    recvd_size += len(data)
                else:
                    data = conn.recv(img_size - recvd_size)
                    recvd_size += len(data)
                # print(recvd_size)
                img_buff += data
            
            img_rec = base64.b64decode(img_buff)
            # print(len(data[6:]))
            with open('./data/test.yuv','wb+') as f:
                f.write(img_rec)
            os.system('ffmpeg -y -s 1920x1080 -i ./data/test.yuv ./data/test.jpg')
            print('recongnition START')
            name = myrecongnition()
            names = ''
            for i in name:
                if i != 'Unknow':
                    names += i + ' ' 
            if names == '':
                names = 'Unknow'
            print('recongnition END')
            print('It is {0}'.format(names))
            print('{0} connection close'.format(addr))
            conn.send(bytes('It is {0}\nConnection closed!'.format(names),'ascii'))
            break
        conn.send(bytes('Hello',"ascii"))
    conn.close()
 
 
if __name__ == '__main__':
    socket_service()

